import pandas as pd
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow.keras import models, layers, losses
from utils.create_log import get_log
from pylab import mpl
mpl.rcParams['font.sans-serif'] = ['SimHei'] #指定默认字体
mpl.rcParams['axes.unicode_minus'] = False #解决保存图像是负号'-'显示为方块的问题


class CovidModel(layers.Layer):
    """
    函数式API构建模型
    """
    def __init__(self, **kwargs):
        super(CovidModel, self).__init__(**kwargs)

    def call(self, x_input, x):
        """

        :param x_input:
        :param x:
        :return:
        """
        return tf.maximum((1 + x) * x_input[:, -1, :], 0.0)

    def get_config(self):
        return super(CovidModel, self).get_config()


class MSPE(losses.Loss):
    """
    自定义损失函数，考虑平方差和预测目标的比值
    """
    def call(self, y_true, y_pred):
        # err_percent = np.power(y_true - y_pred, 2) / (tf.maximum(np.power(y_true, 2), 1e-7))
        err_percent = (y_true - y_pred) ** 2 / tf.maximum(y_true ** 2, 1e-7)
        return tf.reduce_mean(err_percent)

    def get_config(self):
        return super(MSPE, self).get_config()


class Covid19(object):
    def __init__(self):
        self.df = self.get_covid19_data()
        self.WINDOW_SIZE = 8  # 前8日数据预测当日数据
        self.TRAIN_BATCH = 38
        self.OPTIMIZER = tf.keras.optimizers.Adam(learning_rate=0.01)
        # 创建日志
        log_path = '../data/logs/covid_autograph'
        self.logdir = get_log(log_path=log_path)

    @staticmethod
    def get_covid19_data():
        """
        数据获取
        :return:
        """
        csv_file = '../data/covid-19.csv'
        return pd.read_csv(csv_file)

    def get_df_diff(self):
        """
        diff info
        :return:
        """
        df_data = self.df.set_index('date')
        df_diff = df_data.diff(periods=1).dropna()
        return df_diff.reset_index('date')

    def explore_data(self):
        """
        初步探索
        :return:
        """
        self.df.plot(x='date', y=["confirmed_num", "cured_num", "dead_num"])
        plt.xticks(rotation=60)
        plt.legend(['累计确诊人数', '累计治愈人数', '累计死亡人数'])

        df_diff = self.get_df_diff()
        df_diff.plot(x='date', y=["confirmed_num", "cured_num", "dead_num"])
        plt.xticks(rotation=60)
        plt.legend(['累计确诊人数', '累计治愈人数', '累计死亡人数'])
        # df_diff = df_diff.drop('date', axis=1).astype('float32')
        plt.show()

    def prepair(self):
        """
        数据准备
        :return:
        """
        def batch_data(dataset):
            """
            获取前windows_size日数据
            :param dataset:
            :return:
            """
            return dataset.batch(self.WINDOW_SIZE, drop_remainder=True)

        df_diff = self.get_df_diff()
        # 去掉列
        df_diff = df_diff.drop('date', axis=1).astype('float32')
        df_diff_value = df_diff.values

        df_data = tf.data.Dataset.from_tensor_slices(tf.constant(df_diff_value, dtype=tf.float32)).\
            window(self.WINDOW_SIZE, shift=1).flat_map(batch_data)
        df_label = tf.data.Dataset.from_tensor_slices(tf.constant(df_diff_value[self.WINDOW_SIZE:], dtype=tf.float32))

        df_train = tf.data.Dataset.zip((df_data, df_label)).batch(self.TRAIN_BATCH).cache()
        return df_train

    @staticmethod
    def tf_model(shape):
        """
        构建模型
        :return:
        """
        # 清空会话
        tf.keras.backend.clear_session()

        # 函数式API构建模型
        inputs = layers.Input(shape=shape, dtype=tf.float32)
        x = layers.LSTM(shape[1], return_sequences=True, input_shape=shape)(inputs)
        x = layers.LSTM(shape[1], return_sequences=True, input_shape=shape)(x)
        x = layers.LSTM(shape[1], return_sequences=True, input_shape=shape)(x)
        x = layers.LSTM(shape[1], input_shape=shape)(x)
        x = layers.Dense(shape[1])(x)

        x = CovidModel()(inputs, x)
        return models.Model(inputs=[inputs], outputs=[x])

    def train_model(self, model, df_train):
        """
        训练模型
        :param model:
        :param df_train:
        :return:
        """
        model.compile(optimizer=self.OPTIMIZER, loss=MSPE(name='MSPE'))

        tb_callback = tf.keras.callbacks.TensorBoard(log_dir=self.logdir, histogram_freq=1)
        # 若loss在100个epoch后没有提升，learning_rate减半
        lr_callback = tf.keras.callbacks.ReduceLROnPlateau(monitor='loss', factor=0.5, patience=100)
        # 若loss在200个epoch后没有提升，终止训练
        stop_callback = tf.keras.callbacks.EarlyStopping(monitor='loss', patience=200)
        callbacks_list = [tb_callback, lr_callback, stop_callback]

        return model.fit(df_train, epochs=500, callbacks=callbacks_list)

    @staticmethod
    def evalute_model(history, metric):
        """
        评估画图
        :param history:
        :param metric:
        :return:
        """
        train_metrics = history.history[metric]
        epochs = range(1, len(train_metrics) + 1)

        plt.plot(epochs, train_metrics, 'bo--')
        plt.title('训练和验证 ' + metric)
        plt.xlabel('Epochs')
        plt.ylabel(metric)
        plt.legend(['train_' + metric])
        plt.show()

    def predict(self, model, pred_days=100):
        """
        预测
        :return:
        """
        df_diff = self.get_df_diff()
        df_rs = df_diff[["confirmed_num", "cured_num", "dead_num"]].copy()

        for i in range(pred_days):
            arr_predict = model.predict(tf.constant(tf.expand_dims(df_rs.values[-38:, :], axis=0)))
            df_predict = pd.DataFrame(tf.cast(tf.floor(arr_predict), tf.float32).numpy(), columns=df_rs.columns)
            df_rs = df_rs.append(df_predict, ignore_index=True)

        return df_rs

    def main(self):
        # 训练数据
        df_train = self.prepair()
        # 构建模型
        model = self.tf_model(shape=(None, 3))
        # print(model.summary())
        # 训练模型
        history = self.train_model(model=model, df_train=df_train)
        # 评估模型
        self.evalute_model(history=history, metric='loss')
        self.evalute_model(history=history, metric='lr')

        # predict
        df_rs = self.predict(model=model)
        top_10 = df_rs.query('confirmed_num==0').head()
        print(top_10)

        """
        model_path = '../data/models/tf_model_timeseries'
        model.save(model_path, save_format='tf')
        
        model_load = tf.keras.models.load_model(model_path)
        """


def run():
    covid = Covid19()
    # covid.explore_data()
    covid.main()


if __name__ == '__main__':
    run()
